#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Datetime: 2020/11/27 下午2:27
# @Author  : HUANG Xiong
# @Site    : 
# @File    : style_analysis.py
# @Software: PyCharm 

"""
脚本说明：自定义风格分析相关函数
"""
import pandas as pd
import datetime

from quant_researcher.quant.project_tool.math_func.regression_tools import \
    get_rolling_regression_res_by_nav
from quant_researcher.quant.project_tool import time_tool
from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.datasource_fetch.common_data_api import t_trade_date
from quant_researcher.quant.datasource_fetch.fund_api import fund_nav_related
from quant_researcher.quant.datasource_fetch.index_api import index_price_related
from quant_researcher.quant.datasource_fetch.factor_api import factor_return_related
from quant_researcher.quant.datasource_fetch.factor_api.factor_constant import BOND_7_FACTOR, \
    FAMA_5_FACTOR, BARRA_FULL_FACTOR
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.performance_attribution.core_functions.performance_analysis import performance


def style_analysis(target_type, target_code, fund_type, index_list, reg_method, freq,
                   windows, need_bounds, need_constraints, lasso_alpha, ridge_alpha, start_date, end_date):
    """
    输入基金或基金经理代码，指数列表，开始、结束时间，得到回归结果

    :param str target_type: 需要计算的类型，可选1-公募基金，2-公募基金经理，3-私募基金
    :param str target_code: 基金或基金经理的代码
    :param str fund_type: 基金经理管理的基金类型，若target_type=1，3该字段为"00"
    :param list index_list: 指数代码列表，支持传入因子的英文名
    :param int reg_method: 回归算法，1-最小二乘，2-lasso，3-ridge
    :param int freq: 数据频率，1-日频，2-月频。注意，私募基金只有月频数据
    :param int windows: 滚动回归窗口，对target_type=1、2，windows=10表示10天，对target_type=3，windows=10表示10个月
                        为保证回归效果，windows需大于等于10，且大于index_list中index的个数
    :param bool need_bounds: 是否需要给回归系数设置边界，最小二乘设置的边界为(0, 1)，lasso设置的边界为(0, ∞)，
                             ridge无法设置边界条件，此时也需赋值（0或1均可）
    :param bool need_constraints: 是否需要给回归系数设置约束条件，该参数只针对minimize回归，若为True，则系数之和为1，
                                  其他两种回归方法该参数也需赋值（0或1均可）
    :param float lasso_alpha: lasso回归时的alpha参数，若不传，则使用LassoCV寻找最适合的alpha
    :param float ridge_alpha: ridge回归时的alpha参数，若不传，则使用RidgeCV寻找最适合的alpha
    :param str start_date: 计算开始时间，如：2020-05-11
    :param str end_date: 计算截止时间，如：2020-11-11
    :return:
    """
    if reg_method == 1:
        reg_method = 'minimize'
    elif reg_method == 2:
        reg_method = 'lasso'
    else:
        reg_method = 'ridge'

    if target_type in [1, 2]:
        nav_start_date = t_trade_date.date_shifter_by_trade_date(start_date, windows, 'before')
    else:
        nav_start_date = time_tool.date_shifter(start_date, 'months', -windows)
    if target_type == 1:
        f_df = fund_nav_related.get_fund_return(
            fund_code=target_code,
            start_date=nav_start_date,
            end_date=end_date
        )
        f_df = f_df.rename(columns={"end_date": 'trade_date', 'daily_return': 'ret'}).sort_values(by='trade_date')
        f_df = f_df.drop(columns=['fund_code'])
    elif target_type == 2:
        f_df = fund_nav_related.get_manager_return(
            manager_code=target_code,
            start_date=nav_start_date,
            end_date=end_date,
            fund_type=fund_type
        )
        f_df = f_df.rename(columns={"end_date": 'trade_date', 'daily_return': 'ret'}).sort_values(by='trade_date')
        f_df = f_df.drop(columns=['manager_code', 'fund_type'])
    elif target_type == 3:
        conn = db_conn.get_derivative_data_conn()
        f_df = pd.read_sql(f"select month as trade_date, monthly_return as ret "
                           f"from hf_di_fndmonthlyreturn "
                           f"where fund_code = '{target_code}' "
                           f"and month >= '{nav_start_date[:7]}' "
                           f"and month <= '{end_date[:7]}' ", conn)
        conn.close()
        f_df['trade_date'] = f_df['trade_date'].apply(time_tool.get_the_end_of_this_month)
    else:
        raise NotImplementedError

    all_factor_list = BOND_7_FACTOR + FAMA_5_FACTOR + BARRA_FULL_FACTOR
    all_factor_list.append('Default')
    factor_list = []
    for factor in index_list:
        if factor in all_factor_list:
            factor_list.append(factor)
    if target_type in [1, 2]:
        i_df_1 = index_price_related.get_index_return(
            index_code=index_list,
            start_date=nav_start_date,
            end_date=end_date
        )
    else:
        i_df_1 = index_price_related.get_index_return(
            index_code=index_list,
            start_date=nav_start_date,
            end_date=end_date,
            freq='M'
        )
    if i_df_1 is not None:
        i_df_1 = i_df_1.rename(columns={'end_date': 'trade_date'})
        i_df_1 = i_df_1.sort_values(by='trade_date')
        if target_type in [1, 2]:
            i_df_1 = i_df_1.set_index(['trade_date', 'index_code'])['daily_return'].unstack().reset_index()
        else:
            i_df_1 = i_df_1.set_index(['trade_date', 'index_code'])['monthly_return'].unstack().reset_index()

    if target_type in [1, 2]:
        i_df_2 = factor_return_related.get_combined_factor_return(factor_list, nav_start_date, end_date)
    else:
        tmp_date = time_tool.date_shifter(nav_start_date, 'months', -1)
        i_df_2 = factor_return_related.get_combined_factor_return(factor_list, tmp_date, end_date)
    i_df_2 = i_df_2.rename(columns={'end_date': 'trade_date'})

    if len(factor_list) > 0 and target_type == 3:
        i_df_2['trade_date'] = pd.to_datetime(i_df_2['trade_date'])
        i_df_2 = i_df_2.set_index('trade_date')
        i_df_2 = i_df_2.apply(lambda x: performance.aggregate_returns(x, 'monthly')).reset_index()
        i_df_2['trade_date'] = i_df_2.apply(lambda x: datetime.date(int(x['level_0']), int(x['level_1']), 1), axis=1)
        i_df_2['trade_date'] = i_df_2['trade_date'].astype(str)
        i_df_2['trade_date'] = i_df_2['trade_date'].apply(time_tool.get_the_end_of_this_month)
        i_df_2 = i_df_2.drop(columns=['level_0', 'level_1'])
    if i_df_1 is None and len(factor_list) > 0:
        i_df = i_df_2
    elif i_df_1 is not None and len(factor_list) == 0:
        i_df = i_df_1
    elif i_df_1 is not None and len(factor_list) > 0:
        i_df = i_df_1.merge(i_df_2, how='inner', on='trade_date')
    else:
        LOG.error('没有找到输入的index_code的收益率数据，请检查输入')
        return None, None, None
    i_df = i_df.dropna()
    if freq == 2:
        date_list = time_tool.get_date_list(start_date[:7], end_date[:7], front_or_end='end')
        reg_end_list = [x for x in date_list if start_date <= x <= end_date]
        if target_type in [1, 2]:
            reg_start_list = [t_trade_date.date_shifter_by_trade_date(x, windows, 'before') for x in reg_end_list]
        else:
            reg_start_list = [time_tool.date_shifter(x, 'months', -windows) for x in reg_end_list]
        reg_list = list(zip(reg_start_list, reg_end_list))
        res_df = get_rolling_regression_res_by_nav(i_df, f_df, reg_method, True, need_bounds, True, windows,
                                                   reg_list=reg_list, need_constraints=need_constraints,
                                                   lasso_alpha=lasso_alpha, ridge_alpha=ridge_alpha)
    else:
        res_df = get_rolling_regression_res_by_nav(i_df, f_df, reg_method, True, need_bounds, True, windows,
                                                   need_constraints=need_constraints,
                                                   lasso_alpha=lasso_alpha, ridge_alpha=ridge_alpha)

    return res_df, i_df, factor_list


if __name__ == '__main__':
    args = {"target_type": 1,
            "target_code": "110022",
            "fund_type": "00",
            "index_list": ["000929", "000930", "000931", "000936", "000937", "Default", "beta"],
            "reg_method": 2,
            "freq": 1,
            "windows": 60,
            "need_bounds": 1,
            "need_constraints": 1,
            "lasso_alpha": None,
            "ridge_alpha": None,
            "start_date": "2020-11-05",
            "end_date": "2020-11-11"}
    style_analysis(*args.values())
